# -*- coding:utf-8 -*- 

import struct

import gevent
import gevent.event
import gevent.server
import gevent.queue
import gevent.socket

import google.protobuf.service
import message

from . import rpc_pb2 as rpc
from .void_pb2 import Void
from .profile import profile_callmethod, set_profile, get_profile

__all__ = [
    'Error',
    'Controller',
    'Server',
    'Channel',
    'Application',
    'impl',
    'ret',
    'set_profile',
    ]

MAX_ID = 1L << 63
LEN_HEADER = '!I'
LEN_HEADER_SIZE = struct.calcsize(LEN_HEADER)

_id = 0
def _next_id():
    global _id
    _id += 1
    if _id >= MAX_ID:
        _id = 1
    return _id

def impl(service_cls):
    bases = service_cls.__bases__
    for base in bases:
        methods = base.GetDescriptor().methods
        for method in methods:
            m = getattr(service_cls, method.name)
            bm = getattr(base, method.name)
            if m == bm:
                raise RuntimeError('service(%s)\'s method(%s) is not implemented.' \
                        % (service_cls.GetDescriptor().name, method.name))
    return service_cls

def ret(func):
    def _func(self, controller, request, done):
        result = func(self, controller, request)
        done(result)
    return _func

class Poolable(object):
    _spawn = gevent.spawn
    def full(self):
        return False

    def set_spawn(self, spawn):
        if spawn == 'default':
            self.pool = None
            self._spawn = self._spawn
        elif hasattr(spawn, 'spawn'):
            self.pool = spawn
            self._spawn = spawn.spawn
        elif isinstance(spawn, (int, long)):
            from gevent.pool import Pool
            self.pool = Pool(spawn)
            self._spawn = self.pool.spawn
        else:
            self.pool = None
            self._spawn = spawn
        if hasattr(self.pool, 'full'):
            self.full = self.pool.full

class Error(StandardError):
    pass

class Controller(google.protobuf.service.RpcController):
    def __init__(self, channel, *a, **kw):
        google.protobuf.service.RpcController.__init__(self, *a, **kw)
        self.channel = channel

    def Reset(self):
        pass

    def Failed(self):
        pass

    def ErrorText(self):
        pass

    def StartCancel(self):
        pass

    def SetFailed(self, reason):
        pass

    def IsCancelled(self):
        pass

    def NotifyOnCancel(self, callback):
        pass

class Callback(object):
    def __init__(self):
        self.response = None

    def __call__(self, resp):
        self.response = resp

class Sender(object):
    def __init__(self, channel):
        self._queue = gevent.queue.Queue()
        self._channel = channel

    def put(self, buff):
        self._queue.put(buff)

    def loop(self):
        while True:
            buff = self._queue.get()
            packed_buff_len = struct.pack(LEN_HEADER, len(buff))
            try:
                result = self._channel._socket.sendall(packed_buff_len + buff)
            except Exception:
                break
            if result is not None:
                break

class Channel(Poolable, google.protobuf.service.RpcChannel):
    CONNECTED = 'abu.rpc.Channel.CONNECTED'
    DISCONNECTED = 'abu.rpc.Channel.DISCONNECTED'
    def _on_connected(self, host_port):
        pass

    def _on_disconnected(self):
        pass

    def __init__(self, sock, service = None, spawn = 'default', server_side = False):
        google.protobuf.service.RpcChannel.__init__(self)
        self.set_spawn(spawn)
        if not isinstance(sock, gevent.socket.socket):
            s = gevent.socket.socket()
            try:
                s.connect(sock)
            except gevent.socket.error, e:
                raise Error(str(e))
            self._socket = s
        else:
            self._socket = sock
        self._socket.setsockopt(gevent.socket.SOL_TCP, gevent.socket.TCP_NODELAY, 1)
        self._socket.setsockopt(gevent.socket.IPPROTO_TCP, gevent.socket.TCP_NODELAY, 1)

        self._server_side = server_side
        self._closed = False
        self._service = service
        self._pending = {}
        self._sender = Sender(self)
        self._send_let = self._spawn(self._sender.loop)
        self._send_let.link(self.close)
        self._recv_let = self._spawn(self.loop)
        self._recv_let.link(self.close)
        if not self._server_side:
            self._on_connected(sock)
            message.pub(Channel.CONNECTED, self, sock)
    
    def __str__(self):
        try:
            f, t = self.getsockname(), self.getpeername()
            if self._server_side:
                f, t = t, f
            return '<%s.%s: from %s to %s>' % (
                    self.__class__.__module__,
                    self.__class__.__name__, f, t)
        except gevent.socket.error, e:
            return '<%s.%s with socket.error:%s>' % (
                    self.__class__.__module__,
                    self.__class__.__name__, e)

    def __repr__(self):
        return '<%s at %s>' % (str(self), hex(id(self)))

    def getsockname(self):
        return self._socket.getsockname()

    def getpeername(self):
        return self._socket.getpeername()

    def loop(self):
        def recv_bytes(bytes):
            buff = ''
            while len(buff) < bytes:
                t = self._socket.recv(bytes - len(buff))
                if not t:
                    return ''
                buff += t
            return buff

        def recv():
            try:
                buff = recv_bytes(LEN_HEADER_SIZE)
                if not buff:
                    return ''
                packet_len = int(struct.unpack(LEN_HEADER, buff)[0])
                buff = recv_bytes(packet_len)
                return buff
            except gevent.socket.error:
                return ''

        assert self._socket
        assert not self._closed
        while True:
            buff = recv()
            if not buff:
                break
            self._spawn(self._process_packet, buff)
    
    def _process_packet(self, buff):
        packet = rpc.Packet()
        packet.ParseFromString(buff)
        if packet.type == rpc.REQUEST:
            self._process_request(packet.identify, packet.serialized)
        elif packet.type == rpc.RESPONSE:
            self._process_response(packet.identify, packet.serialized)
        elif packet.type == rpc.ERROR:
            error = rpc.Error()
            error.ParseFromString(packet.serialized)
            self._process_error(packet.identify, error.info)
        else:
            self._process_error(packet.identify, 'unsupported packet:%s' % repr(packet))

    def _send_error(self, identify, error_info):
        error = rpc.Error()
        error.info = error_info
        packet = rpc.Packet()
        packet.type = rpc.ERROR
        packet.identify = identify
        packet.serialized = error.SerializeToString()
        self._send(packet.SerializeToString())

    def _process_request(self, identify, buff):
        req = rpc.Request()
        req.ParseFromString(buff)
        
        service = self._service

        method = service.GetDescriptor().FindMethodByName(req.method)

        if not method:
            self._send_error(identify, u'method(%s) is not found.' % (req.method))
            return

        if req.HasField('request'):
            request = service.GetRequestClass(method)()
            request.ParseFromString(req.request)
        else:
            request = None
        controller = Controller(self)
        callback = Callback()
        try:
            service.CallMethod(method, controller, request, callback)
        except Exception, e:
            self._send_error(identify, u'some error occurred while call method(%s.%s): %s' % (
                service.GetDescriptor().name, req.method, e))
            raise

        resp = rpc.Response()
        if callback.response:
            resp.response = callback.response.SerializeToString()
        else:
            resp_cls = service.GetResponseClass(method)
            if resp_cls == Void:
                return
            else:
                e = u'method(%s) return None.' % (req.method)
                self._send_error(identify, e)
                return

        packet = rpc.Packet()
        packet.type = rpc.RESPONSE
        packet.identify = identify
        packet.serialized = resp.SerializeToString()
        self._send(packet.SerializeToString())

    def _process_response(self, identify, buff):
        try:
            result = self._pending[identify]
        except KeyError:
            return
        resp = rpc.Response()
        resp.ParseFromString(buff)
        result.set(resp.response)

    def _process_error(self, identify, error_info):
        try:
            result = self._pending[identify]
        except KeyError:
            return
        result.set_exception(Error(error_info))

    def CallMethod(self, mdesc, controller, request, response_class, done = None):
        req_id = _next_id()
        assert req_id not in self._pending
        if response_class != Void:
            result = gevent.event.AsyncResult()
            self._pending[req_id] = result

        try:
            req = rpc.Request()
            req.method = mdesc.name
            req.request = request.SerializeToString()

            packet = rpc.Packet()
            packet.type = rpc.REQUEST
            packet.identify = req_id
            packet.serialized = req.SerializeToString()
            self._send(packet.SerializeToString())
            
            if response_class == Void:
                return

            resp = response_class()
            resp_buff = result.get()
            resp.ParseFromString(resp_buff)
            done(resp)
        finally:
            try:
                self._pending.pop(req_id)
            except KeyError:
                pass

    def _send(self, buff):
        if self._closed:
            return
        self._sender.put(buff)

    def close(self, t = None):
        if self._closed:
            return
        self._closed = True
        self._spawn(self._real_close)
        if not self._server_side:
            self._on_disconnected()
            message.pub(Channel.DISCONNECTED, self)

    def _real_close(self):
        if self._socket:
            self._socket.close()
            self._socket = None
        if self._send_let:
            self._send_let.kill()
            self._send_let = None
        if self._recv_let:
            self._recv_let.kill()
            self._recv_let = None
        self._sender = None

class Proxy(object):
    def __init__(self, stub):
        self._stub = stub

    def __dir__(self):
        return dir(self._stub)

    def __getattr__(self, name):
        def call(method, request):
            controller = Controller(None)
            callback = Callback()
            method(controller, request, callback)
            return callback.response

        if get_profile():
            call = profile_callmethod(
                    '.'.join((self._stub.GetDescriptor().name, name)))(call)

        req_cls = self._stub.GetRequestClass(
                self._stub.GetDescriptor().FindMethodByName(name))

        def rpc_method(*a, **kw):
            if (not a) and (not kw):
                req = Void()
                return call(getattr(self._stub, name), req)
            if a:
                if len(a) != 1:
                    raise TypeError('%s() takes at most 1 argument (%d given)' % (name, len(a)))
                return call(getattr(self._stub, name), a[0])
            if kw:
                req = req_cls(**kw)
                return call(getattr(self._stub, name), req)
            raise RuntimeError('Some error occurred.')

        return rpc_method

class Server(gevent.server.StreamServer):
    NEW_CONNECTION = 'abu.rpc.Server.NEW_CONNECTION'
    LOST_CONNECTION = 'abu.rpc.Server.LOST_CONNECTION'

    def _on_new_connection(self, channel):
        pass

    def _on_lost_connection(self, channel):
        pass

    def __init__(self, listener, service, backlog = None, spawn='default', **ssl_args):
        gevent.server.StreamServer.__init__(self,
                listener,
                self._handler,
                backlog,
                spawn,
                **ssl_args)
        self._service = service

    def __str__(self):
        return '<%s.%s:%s>' % (
                self.__class__.__module__,
                self.__class__.__name__,
                str(self.address))

    def __repr__(self):
        return '<%s at %s>' % (str(self), hex(id(self)))

    def _handler(self, sock, addr):
        channel = Channel(sock, self._service, self._spawn, server_side = True)
        self._on_new_connection(channel)
        message.pub(Server.NEW_CONNECTION, self, channel)
        channel._recv_let.join()
        self._on_lost_connection(channel)
        message.pub(Server.LOST_CONNECTION, self, channel)

class Application(Poolable):
    def __init__(self, *servers):
        self._servers = servers

    def run(self):
        jobs = [self._spawn(server.serve_forever) for server in self._servers]
        gevent.joinall(jobs)

